# 计算超左时间窗惩罚成本（提前服务惩罚成本），参数为测试用distance_list
from functools import partial

import numpy as np
import torch


def vrtpw_cost(cars_list, data, device=None):
    WT = np.array([data.QT, data.FT]).T.astype(np.float32)
    ST = data.ST
    route_dist = data.route_dist
    time_distance_list = data.route_time
    start_time = data.QT - data.route_time[0, 0, :]

    myfun = partial(gpu_vrptw_cost,
                    WT=torch.from_numpy(WT).to(device),
                    ST=torch.from_numpy(ST).to(device),
                    route_dist=torch.from_numpy(route_dist).to(device),
                    time_distance_list=torch.from_numpy(time_distance_list).to(device),
                    start_time=torch.from_numpy(start_time).to(device)
                    )
    ans = myfun(cars_list)

    return ans


def raw_vrptw_cost(cars_list, WT, ST, route_dist, time_distance_list, start_time):
    cars_list_before = cars_list

    cur, time_dist, router_cost = np.zeros((3, cars_list.shape[0]), dtype=np.float32)

    n = cars_list_before.shape[-1]
    for i in range(n):
        car_before = cars_list_before[:, i]
        car_after = cars_list_before[:, (i + 1) % n]

        arrival_car = np.where(car_before == 0, car_after, 0 * car_after)
        cur = np.where(arrival_car != 0, start_time[arrival_car], cur)

        cur_step = np.clip(cur / 30 - 1, 0, time_distance_list.shape[0] - 1).astype(np.int16)
        route_index = (
            cur_step,
            car_before,
            car_after
        )

        cur += time_distance_list[route_index]
        router_cost += route_dist[route_index]

        cur_WT = WT[car_after]
        cur = np.maximum(cur, cur_WT[:, 0])
        time_dist += np.maximum(0, cur - cur_WT[:, 1])
        cur += ST[car_after]

    return router_cost + time_dist * 1e5


def gpu_vrptw_cost(cars_list, WT, ST, route_dist, time_distance_list, start_time):
    n = cars_list.shape[-1]

    cur, Pun_cost, router_cost = torch.zeros((3, cars_list.shape[0]), dtype=torch.float32,
                                             device=cars_list.device)

    for i in range(n):
        car_before = cars_list[:, i].long()
        car_after = cars_list[:, (i + 1) % n].long()

        arrival_car = torch.where(car_before == 0, car_after, 0 * car_after)
        cur = torch.where(arrival_car != 0, start_time[arrival_car], cur)

        cur_time_step = torch.clamp(cur / 30 - 1, 0, time_distance_list.shape[0] - 1)

        indx = (
            cur_time_step.long(),
            car_before,
            car_after
        )

        cur += time_distance_list[indx]
        router_cost += route_dist[indx]

        cur_WT = WT[car_after]
        cur = torch.maximum(cur, cur_WT[:, 0])
        Pun_cost += torch.relu(cur - cur_WT[:, 1])
        cur += ST[car_after]

    time_dist = Pun_cost
    return (router_cost + time_dist * 1e5)


def faster_gpu_vrptw_cost(cars_list, WT, ST, route_dist, time_distance_list, start_time):
    n = cars_list.shape[-1]

    sites = start_time.shape[0]
    cars = cars_list.shape[-1] - sites + 1
    max_l = sites * 3 // cars

    cars_list_before = torch.zeros((cars_list.shape[0], cars, max_l), dtype=torch.int16, device=cars_list.device)
    max_len, cur_cars, cur_site = torch.zeros((3, cars_list.shape[0]), dtype=torch.int64, device=cars_list.device)
    cur_cars -= 1
    indx = torch.arange(cars_list.shape[0], device=cars_list.device)
    for i in range(n):
        arr = cars_list[:, i] == 0
        cur_cars += arr
        cur_site[arr] = 0
        cars_list_before[indx, cur_cars, cur_site] = cars_list[:, i]
        cur_site += 1
        max_len = torch.maximum(max_len, cur_site)

    n = torch.max(max_len).item()

    cars_list_before = cars_list_before.reshape((-1, max_l))

    cur, time_dist, router_cost = torch.zeros((3, cars_list_before.shape[0]), dtype=torch.float32,
                                              device=cars_list_before.device)

    for i in range(n):
        car_before = cars_list_before[:, i].long()
        car_after = cars_list_before[:, (i + 1) % n].long()

        arrival_car = torch.where(car_before == 0, car_after, 0 * car_after)
        cur = torch.where(arrival_car != 0, start_time[arrival_car], cur)

        cur_time_step = torch.clamp(cur / 30 - 1, 0, time_distance_list.shape[0] - 1)

        indx = (
            cur_time_step.long(),
            car_before,
            car_after
        )

        cur += time_distance_list[indx]
        router_cost += route_dist[indx]

        cur_WT = WT[car_after]
        cur = torch.maximum(cur, cur_WT[:, 0])
        time_dist += torch.relu(cur - cur_WT[:, 1])
        cur += ST[car_after]

    time_dist = torch.sum(time_dist.reshape((cars_list.shape[0], -1)), -1)
    router_cost = torch.sum(router_cost.reshape((cars_list.shape[0], -1)), -1)
    return (router_cost + time_dist * 1e5)
